package main

import (
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/go-rod/rod"
	"github.com/gorilla/websocket"
)

func main() {
	var pic []byte

	page := rod.New().MustConnect().MustPage("https://www.iqiyi.com/")
	page.MustWaitLoad()

	go func() {
		for {
			pic = page.MustElement("#block-B").MustScreenshot("")

			time.Sleep(time.Second)
		}
	}()

	r := gin.Default()
	r.Static("/assets", "./web")
	r.GET("/test", func(c *gin.Context) {
		c.JSON(http.StatusOK, []string{"123", "321"})
	})
	r.GET("/ws", func(c *gin.Context) {
		upGrader := websocket.Upgrader{
			CheckOrigin: func(r *http.Request) bool {
				return true
			},
		}
		// 升级get请求为webSocket协议
		ws, err := upGrader.Upgrade(c.Writer, c.Request, nil)
		if err != nil {
			return
		}
		defer ws.Close()

		for {
			mt, message, err := ws.ReadMessage()
			if err != nil {
				break
			}

			str := string(message)

			if str == "ping" {
				message = []byte("pong")
			}

			if str == "pic" {
				err = ws.WriteMessage(websocket.BinaryMessage, pic)
				if err != nil {
					break
				}

				continue
			}
			// 写入ws数据
			err = ws.WriteMessage(mt, message)
			if err != nil {
				break
			}
		}
	})

	_ = r.Run(":8080")
}
